package layers

import (
	"math"
)

type paddingLayer struct {
	padding int
}

func NewPaddingLayer(padding int) *paddingLayer {
	return &paddingLayer{padding: padding}
}

func (pl *paddingLayer) Forward(input []float64) []float64 {
	//log.Println("start padding")
	input2d := ConvertTensor1to2(input, uint(math.Sqrt(float64(len(input)))))

	output := make([][]float64, len(input2d)+pl.padding*2)
	for i := range output {
		output[i] = make([]float64, len(input2d[0])+pl.padding*2)
	}

	//log.Println("Padding:", len(input2d), "x", len(input2d[0]), "->", len(output), "x", len(output[0]))
	for i := 2; i < len(output[0])-2; i++ {
		tmp := make([]float64, 0)
		tmp = append(tmp, 0, 0)
		tmp = append(tmp, input2d[i-2]...)
		tmp = append(tmp, 0, 0)
		output[i] = tmp
	}
	//log.Println("end padding")

	return ConvertTensor2to1(output)

}

func (pl *paddingLayer) BackProp(output []float64) []float64 {
	//log.Println("back padding")
	output2d := ConvertTensor1to2(output, uint(math.Sqrt(float64(len(output)))))
	res := make([][]float64, 0)
	for _, v := range output2d {
		res = append(res, v[2:len(v)-2])
	}
	//log.Println("end back padding")
	return ConvertTensor2to1(res)
}
